# encoding: utf-8


"""

@author: tongzhenguo

@time: 2020/4/8 上午10:51

@desc:

封装内容池条件过滤(白名单、黑名单)

"""

from pandas import DataFrame
from data.mongo_db import MongoDB


def dataframe_in_set_pool(df: DataFrame, col, in_set, not_in_set):
    """
    mongo根据某个字段条件进行过滤,支持in逻辑运算
    :param df: 输入数据dataframe
    :param col: 逻辑in、not in运算的字段名
    :param in_set: 白名单
    :param not_in_set: 黑名单
    :return: 满足in(白名单)且not in运算(黑名单)的内容池
    """
    if in_set:
        df = df[df[col].isin(in_set)]
    if not_in_set:
        df = df[~df[col].isin(not_in_set)]
    return df


def mongo_in_set_pool(mongo_db: MongoDB, col, in_set, not_in_set, collection_name, column_list, query_dict):
    """
    mongo根据某个字段条件进行过滤,支持in逻辑运算
    :param col: 逻辑in、not in运算的字段名
    :param mongo_db: mongo_db客户端
    :param in_set: 白名单
    :param not_in_set: 黑名单
    :param collection_name: mongo集合名
    :param column_list: mongo拉取字段列表
    :param query_dict: mongo查询
    :return: 满足in(白名单)且not in运算(黑名单)的内容池
    """
    typed_df = mongo_db.fetch_data(column_list, collection_name, query_dict)
    typed_df = dataframe_in_set_pool(typed_df, col, in_set, not_in_set)
    return typed_df


if __name__ == "__main__":
    import pandas as pd

    data = pd.DataFrame([[1, "a"]], columns=['id', 'tag'])
    dataframe_in_set_pool(data, "tag", ["a"], None)
